{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "using MLDatasets, Plots, LinearAlgebra, ProgressMeter, Statistics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "using Flux\n",
    "using Flux: binarycrossentropy\n",
    "using Flux: params #Used for automatic differenation (but in Project replace auto-diff with explicit gradient)\n",
    "using Statistics, Random, StatsBase\n",
    "using Flux: onehotbatch, crossentropy, update!\n",
    "Random.seed!(0);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_x, train_y = MNIST(split=:train)[:];\n",
    "test_x,  test_y  = MNIST(split=:test)[:];"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(60000, 785)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X_total = hcat(ones(60000), hcat([float.(vec(train_x[:,:,i])) for i in 1:60000]...)');\n",
    "size(X_total)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10-element Vector{Tuple{Int64, Int64}}:\n",
       " (5923, 785)\n",
       " (6742, 785)\n",
       " (5958, 785)\n",
       " (6131, 785)\n",
       " (5842, 785)\n",
       " (5421, 785)\n",
       " (5918, 785)\n",
       " (6265, 785)\n",
       " (5851, 785)\n",
       " (5949, 785)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X_per_digit = [X_total[train_y .== i, :] for i in 0:9]\n",
    "size.(X_per_digit)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# (1) Linear One vs. Rest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "accuracy_one_vs_rest = 0.8603\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "10×10 Matrix{Int64}:\n",
       " 944     0   18    4    0   23   18    5   14   15\n",
       "   0  1107   54   17   22   18   10   40   46   11\n",
       "   1     2  813   23    6    3    9   16   11    2\n",
       "   2     2   26  880    1   72    0    6   30   17\n",
       "   2     3   15    5  881   24   22   26   27   80\n",
       "   7     1    0   17    5  659   17    0   40    1\n",
       "  14     5   42    9   10   23  875    1   15    1\n",
       "   2     1   22   21    2   14    0  884   12   77\n",
       "   7    14   37   22   11   39    7    0  759    4\n",
       "   1     0    5   12   44   17    0   50   20  801"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pinv_total = pinv(X_total)\n",
    "\n",
    "\"\"\"\n",
    "Make a classifier for the `digit_pos`.\n",
    "\"\"\"\n",
    "function make_one_vs_one_pred(digit_pos)\n",
    "    Y_one_vs_rest = 2(train_y .== digit_pos) .- 1\n",
    "    \n",
    "    #beta is the set of weights for the specific digit...\n",
    "    β_one_vs_rest = pinv_total*Y_one_vs_rest\n",
    "    pred_one_vs_rest(img) = vcat(1,vec(img))'β_one_vs_rest\n",
    "    return pred_one_vs_rest\n",
    "end\n",
    "\n",
    "preds_one_vs_rest = [make_one_vs_one_pred(i) for i in 0:9]\n",
    "\n",
    "#Here is the \"one vs. rest\"\n",
    "predict_one_vs_rest(img) = argmax([preds_one_vs_rest[i+1](img) for i in 0:9]) - 1\n",
    "\n",
    "accuracy_one_vs_rest = mean([predict_one_vs_rest(test_x[:,:,i]) == test_y[i] for i in 1:10000]) \n",
    "\n",
    "@show accuracy_one_vs_rest;\n",
    "\n",
    "function conf_matrix(pred_fun) \n",
    "    predictions = [predict_one_vs_rest(test_x[:,:,i]) for i in 1:10000]\n",
    "    confusionMatrix = [sum((predictions .== i) .& (test_y .== j)) for i in 0:9, j in 0:9]\n",
    "end\n",
    "conf_matrix(predict_one_vs_rest)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# (2) Linear One vs. One"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[32mProgress: 100%|█████████████████████████████████████████| Time: 0:00:08\u001b[39m\n",
      "\u001b[32mProgress: 100%|█████████████████████████████████████████| Time: 0:00:08\u001b[39m\n",
      "\u001b[32mProgress: 100%|█████████████████████████████████████████| Time: 0:00:07\u001b[39m\n",
      "\u001b[32mProgress: 100%|█████████████████████████████████████████| Time: 0:00:07\u001b[39m\n",
      "\u001b[32mProgress: 100%|█████████████████████████████████████████| Time: 0:00:07\u001b[39m\n",
      "\u001b[32mProgress: 100%|█████████████████████████████████████████| Time: 0:00:08\u001b[39m\n",
      "\u001b[32mProgress: 100%|█████████████████████████████████████████| Time: 0:00:09\u001b[39m\n",
      "\u001b[32mProgress: 100%|█████████████████████████████████████████| Time: 0:00:09\u001b[39m\n",
      "\u001b[32mProgress: 100%|█████████████████████████████████████████| Time: 0:00:09\u001b[39m\n",
      "\u001b[32mProgress: 100%|█████████████████████████████████████████| Time: 0:00:09\u001b[39m\n",
      "\u001b[32mProgress: 100%|█████████████████████████████████████████| Time: 0:01:26\u001b[39m\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "accuracy_one_vs_one_sign = 0.9297\n"
     ]
    }
   ],
   "source": [
    "function make_one_vs_one_pred(digit_pos, digit_neg)\n",
    "    X_one_vs_one = vcat(X_per_digit[digit_pos+1],X_per_digit[digit_neg+1])\n",
    "    Y_one_vs_one = vcat(ones(size(X_per_digit[digit_pos+1])[1]), -ones(size(X_per_digit[digit_neg+1])[1]))\n",
    "    β_one_vs_one = pinv(X_one_vs_one)*Y_one_vs_one\n",
    "    pred_one_vs_one(img) = vcat(1,vec(img))'β_one_vs_one\n",
    "    return pred_one_vs_one\n",
    "end\n",
    "\n",
    "preds_one_vs_one = Dict()\n",
    "@showprogress for i in 0:9\n",
    "    @showprogress for j in 0:9\n",
    "        i == j && continue\n",
    "        preds_one_vs_one[(i,j)] = make_one_vs_one_pred(i,j)\n",
    "    end\n",
    "end\n",
    "\n",
    "predict_one_vs_one_sign(img) = argmax([sum([sign(preds_one_vs_one[(i,j)](img)) for j in setdiff(0:9,i)]) for i in 0:9])-1\n",
    "\n",
    "accuracy_one_vs_one_sign = mean([predict_one_vs_one_sign(test_x[:,:,i]) == test_y[i] for i in 1:10000]) \n",
    "\n",
    "@show accuracy_one_vs_one_sign;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10×10 Matrix{Int64}:\n",
       " 944     0   18    4    0   23   18    5   14   15\n",
       "   0  1107   54   17   22   18   10   40   46   11\n",
       "   1     2  813   23    6    3    9   16   11    2\n",
       "   2     2   26  880    1   72    0    6   30   17\n",
       "   2     3   15    5  881   24   22   26   27   80\n",
       "   7     1    0   17    5  659   17    0   40    1\n",
       "  14     5   42    9   10   23  875    1   15    1\n",
       "   2     1   22   21    2   14    0  884   12   77\n",
       "   7    14   37   22   11   39    7    0  759    4\n",
       "   1     0    5   12   44   17    0   50   20  801"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conf_matrix(predict_one_vs_one_sign)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# (3) Logistic regression one vs. rest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_y_digit_one_vs_rest = [train_y .== k for k in 0:9];"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "train_logistic (generic function with 1 method)"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sig(x) = 1/(1+float(MathConstants.e)^-x)\n",
    "logistic_predict(img_vec, w) = sig.(w'*img_vec);\n",
    "\n",
    "function train_logistic(data_x, data_y, train_number;\n",
    "        num_epochs = 200, \n",
    "        mini_batch_size = 2000, \n",
    "        η = 0.02)\n",
    "    println(\"starting $train_number\")\n",
    "    \n",
    "    w = randn(28*28+1)\n",
    "    loss(x, y) = binarycrossentropy(logistic_predict(x, w), y);\n",
    "    loss_value = 0.0\n",
    "    opt = ADAM(η)\n",
    "    n = size(data_x)[1]\n",
    "    for epoch_num in 1:num_epochs\n",
    "        for batch in Iterators.partition(1:n, mini_batch_size)\n",
    "            gs = gradient(()->loss( data_x'[:,batch], \n",
    "                                    data_y[batch]'), \n",
    "                                    params(w))\n",
    "            for p in (w,)\n",
    "                update!(opt, p, gs[p])\n",
    "            end\n",
    "\n",
    "        end\n",
    "        print(\".\")\n",
    "    end\n",
    "    println()\n",
    "    return w\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "starting 0\n",
      "........................................................................................................................................................................................................\n",
      "starting 1\n",
      "........................................................................................................................................................................................................\n",
      "starting 2\n",
      "........................................................................................................................................................................................................\n",
      "starting 3\n",
      "........................................................................................................................................................................................................\n",
      "starting 4\n",
      "........................................................................................................................................................................................................\n",
      "starting 5\n",
      "........................................................................................................................................................................................................\n",
      "starting 6\n",
      "........................................................................................................................................................................................................\n",
      "starting 7\n",
      "........................................................................................................................................................................................................\n",
      "starting 8\n",
      "........................................................................................................................................................................................................\n",
      "starting 9\n",
      "........................................................................................................................................................................................................\n"
     ]
    }
   ],
   "source": [
    "logistic_one_rest_models = [train_logistic(X_total, train_y_digit_one_vs_rest[k+1], k) for k in 0:9];"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "accuracy_one_vs_rest_log = 0.9174\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "10×10 Matrix{Int64}:\n",
       " 944     0   18    4    0   23   18    5   14   15\n",
       "   0  1107   54   17   22   18   10   40   46   11\n",
       "   1     2  813   23    6    3    9   16   11    2\n",
       "   2     2   26  880    1   72    0    6   30   17\n",
       "   2     3   15    5  881   24   22   26   27   80\n",
       "   7     1    0   17    5  659   17    0   40    1\n",
       "  14     5   42    9   10   23  875    1   15    1\n",
       "   2     1   22   21    2   14    0  884   12   77\n",
       "   7    14   37   22   11   39    7    0  759    4\n",
       "   1     0    5   12   44   17    0   50   20  801"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function make_one_vs_one_pred_logistic(digit_pos)\n",
    "    Y_one_vs_rest = 2(train_y .== digit_pos) .- 1\n",
    "    w = logistic_one_rest_models[digit_pos]\n",
    "    predict_f(img) = logistic_predict(vcat(1,vec(img)), w)\n",
    "    return predict_f\n",
    "end\n",
    "\n",
    "logistic_pred = [make_one_vs_one_pred_logistic(k) for k in 1:10];\n",
    "pred_img_log_one_rest(img) = argmax([logistic_pred[k](img) for k in 1:10])-1\n",
    "\n",
    "accuracy_one_vs_rest_log = mean([pred_img_log_one_rest(test_x[:,:,i]) == test_y[i] for i in 1:10000]) \n",
    "\n",
    "@show accuracy_one_vs_rest_log;\n",
    "\n",
    "conf_matrix(pred_img_log_one_rest)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# (4) Logistic regression one vs. one"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "starting (0, 1)\n",
      "........................................................................................................................................................................................................\n",
      "starting (0, 2)\n",
      "........................................................................................................................................................................................................\n",
      "starting (0, 3)\n",
      "........................................................................................................................................................................................................\n",
      "starting (0, 4)\n",
      "........................................................................................................................................................................................................\n",
      "starting (0, 5)\n",
      "........................................................................................................................................................................................................\n",
      "starting (0, 6)\n",
      "........................................................................................................................................................................................................\n",
      "starting (0, 7)\n",
      "........................................................................................................................................................................................................\n",
      "starting (0, 8)\n",
      "........................................................................................................................................................................................................\n",
      "starting (0, 9)\n",
      "........................................................................................................................................................................................................\n",
      "starting (1, 0)\n",
      "........................................................................................................................................................................................................\n",
      "starting (1, 2)\n",
      "........................................................................................................................................................................................................\n",
      "starting (1, 3)\n",
      "........................................................................................................................................................................................................\n",
      "starting (1, 4)\n",
      "........................................................................................................................................................................................................\n",
      "starting (1, 5)\n",
      "........................................................................................................................................................................................................\n",
      "starting (1, 6)\n",
      "........................................................................................................................................................................................................\n",
      "starting (1, 7)\n",
      "........................................................................................................................................................................................................\n",
      "starting (1, 8)\n",
      "........................................................................................................................................................................................................\n",
      "starting (1, 9)\n",
      "........................................................................................................................................................................................................\n",
      "starting (2, 0)\n",
      "........................................................................................................................................................................................................\n",
      "starting (2, 1)\n",
      "........................................................................................................................................................................................................\n",
      "starting (2, 3)\n",
      "........................................................................................................................................................................................................\n",
      "starting (2, 4)\n",
      "........................................................................................................................................................................................................\n",
      "starting (2, 5)\n",
      "........................................................................................................................................................................................................\n",
      "starting (2, 6)\n",
      "........................................................................................................................................................................................................\n",
      "starting (2, 7)\n",
      "........................................................................................................................................................................................................\n",
      "starting (2, 8)\n",
      "........................................................................................................................................................................................................\n",
      "starting (2, 9)\n",
      "........................................................................................................................................................................................................\n",
      "starting (3, 0)\n",
      "........................................................................................................................................................................................................\n",
      "starting (3, 1)\n",
      "........................................................................................................................................................................................................\n",
      "starting (3, 2)\n",
      "........................................................................................................................................................................................................\n",
      "starting (3, 4)\n",
      "........................................................................................................................................................................................................\n",
      "starting (3, 5)\n",
      "........................................................................................................................................................................................................\n",
      "starting (3, 6)\n",
      "........................................................................................................................................................................................................\n",
      "starting (3, 7)\n",
      "........................................................................................................................................................................................................\n",
      "starting (3, 8)\n",
      "........................................................................................................................................................................................................\n",
      "starting (3, 9)\n",
      "........................................................................................................................................................................................................\n",
      "starting (4, 0)\n",
      "........................................................................................................................................................................................................\n",
      "starting (4, 1)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "........................................................................................................................................................................................................\n",
      "starting (4, 2)\n",
      "........................................................................................................................................................................................................\n",
      "starting (4, 3)\n",
      "........................................................................................................................................................................................................\n",
      "starting (4, 5)\n",
      "........................................................................................................................................................................................................\n",
      "starting (4, 6)\n",
      "........................................................................................................................................................................................................\n",
      "starting (4, 7)\n",
      "........................................................................................................................................................................................................\n",
      "starting (4, 8)\n",
      "........................................................................................................................................................................................................\n",
      "starting (4, 9)\n",
      "........................................................................................................................................................................................................\n",
      "starting (5, 0)\n",
      "........................................................................................................................................................................................................\n",
      "starting (5, 1)\n",
      "........................................................................................................................................................................................................\n",
      "starting (5, 2)\n",
      "........................................................................................................................................................................................................\n",
      "starting (5, 3)\n",
      "........................................................................................................................................................................................................\n",
      "starting (5, 4)\n",
      "........................................................................................................................................................................................................\n",
      "starting (5, 6)\n",
      "........................................................................................................................................................................................................\n",
      "starting (5, 7)\n",
      "........................................................................................................................................................................................................\n",
      "starting (5, 8)\n",
      "........................................................................................................................................................................................................\n",
      "starting (5, 9)\n",
      "........................................................................................................................................................................................................\n",
      "starting (6, 0)\n",
      "........................................................................................................................................................................................................\n",
      "starting (6, 1)\n",
      "........................................................................................................................................................................................................\n",
      "starting (6, 2)\n",
      "........................................................................................................................................................................................................\n",
      "starting (6, 3)\n",
      "........................................................................................................................................................................................................\n",
      "starting (6, 4)\n",
      "........................................................................................................................................................................................................\n",
      "starting (6, 5)\n",
      "........................................................................................................................................................................................................\n",
      "starting (6, 7)\n",
      "........................................................................................................................................................................................................\n",
      "starting (6, 8)\n",
      "........................................................................................................................................................................................................\n",
      "starting (6, 9)\n",
      "........................................................................................................................................................................................................\n",
      "starting (7, 0)\n",
      "........................................................................................................................................................................................................\n",
      "starting (7, 1)\n",
      "........................................................................................................................................................................................................\n",
      "starting (7, 2)\n",
      "........................................................................................................................................................................................................\n",
      "starting (7, 3)\n",
      "........................................................................................................................................................................................................\n",
      "starting (7, 4)\n",
      "........................................................................................................................................................................................................\n",
      "starting (7, 5)\n",
      "........................................................................................................................................................................................................\n",
      "starting (7, 6)\n",
      "........................................................................................................................................................................................................\n",
      "starting (7, 8)\n",
      "........................................................................................................................................................................................................\n",
      "starting (7, 9)\n",
      "........................................................................................................................................................................................................\n",
      "starting (8, 0)\n",
      "........................................................................................................................................................................................................\n",
      "starting (8, 1)\n",
      "........................................................................................................................................................................................................\n",
      "starting (8, 2)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "........................................................................................................................................................................................................\n",
      "starting (8, 3)\n",
      "........................................................................................................................................................................................................\n",
      "starting (8, 4)\n",
      "........................................................................................................................................................................................................\n",
      "starting (8, 5)\n",
      "........................................................................................................................................................................................................\n",
      "starting (8, 6)\n",
      "........................................................................................................................................................................................................\n",
      "starting (8, 7)\n",
      "........................................................................................................................................................................................................\n",
      "starting (8, 9)\n",
      "........................................................................................................................................................................................................\n",
      "starting (9, 0)\n",
      "........................................................................................................................................................................................................\n",
      "starting (9, 1)\n",
      "........................................................................................................................................................................................................\n",
      "starting (9, 2)\n",
      "........................................................................................................................................................................................................\n",
      "starting (9, 3)\n",
      "........................................................................................................................................................................................................\n",
      "starting (9, 4)\n",
      "........................................................................................................................................................................................................\n",
      "starting (9, 5)\n",
      "........................................................................................................................................................................................................\n",
      "starting (9, 6)\n",
      "........................................................................................................................................................................................................\n",
      "starting (9, 7)\n",
      "........................................................................................................................................................................................................\n",
      "starting (9, 8)\n",
      "........................................................................................................................................................................................................\n"
     ]
    }
   ],
   "source": [
    "function make_one_vs_one_pred_x_y(digit_pos, digit_neg)\n",
    "    X_one_vs_one = vcat(X_per_digit[digit_pos+1],X_per_digit[digit_neg+1])\n",
    "    Y_one_vs_one = vcat(ones(size(X_per_digit[digit_pos+1])[1]), zeros(size(X_per_digit[digit_neg+1])[1]))\n",
    "    return X_one_vs_one, Y_one_vs_one\n",
    "end\n",
    "\n",
    "preds_one_vs_one_log = Dict()\n",
    "for i in 0:9\n",
    "    for j in 0:9\n",
    "        i == j && continue\n",
    "        x, y = make_one_vs_one_pred_x_y(i,j)\n",
    "        preds_one_vs_one_log[(i,j)] = train_logistic(x, y, (i,j))\n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "accuracy_one_vs_rest_log = 0.9321\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "10×10 Matrix{Int64}:\n",
       " 944     0   18    4    0   23   18    5   14   15\n",
       "   0  1107   54   17   22   18   10   40   46   11\n",
       "   1     2  813   23    6    3    9   16   11    2\n",
       "   2     2   26  880    1   72    0    6   30   17\n",
       "   2     3   15    5  881   24   22   26   27   80\n",
       "   7     1    0   17    5  659   17    0   40    1\n",
       "  14     5   42    9   10   23  875    1   15    1\n",
       "   2     1   22   21    2   14    0  884   12   77\n",
       "   7    14   37   22   11   39    7    0  759    4\n",
       "   1     0    5   12   44   17    0   50   20  801"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "logistic_predict_one_vs_one(i,j,img) = logistic_predict(vcat(1,vec(img)), \n",
    "                                                preds_one_vs_one_log[(i,j)]) - 0.5\n",
    "\n",
    "predict_one_vs_one_sign_log(img) = argmax([sum([sign(logistic_predict_one_vs_one(i,j,img)) for j in setdiff(0:9,i)]) for i in 0:9])-1\n",
    "\n",
    "accuracy_one_vs_rest_log = mean([predict_one_vs_one_sign_log(test_x[:,:,i]) == test_y[i] for i in 1:10000])\n",
    "\n",
    "@show accuracy_one_vs_rest_log\n",
    "conf_matrix(predict_one_vs_one_sign_log)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# (5) Multiclass classifier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_test = length(test_y);\n",
    "n_train = length(train_y);\n",
    "X_test = vcat([vec(test_x[:,:,k])' for k in 1:n_test]...);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "train_softmax_logistic (generic function with 1 method)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "logistic_softmax_predict(img_vec, W) = softmax(W*img_vec)\n",
    "\n",
    "function train_softmax_logistic(;mini_batch_size = 2000, num_epochs = 200)\n",
    "    \n",
    "    #Initilize parameters\n",
    "    W = randn(10,28*28+1)\n",
    "\n",
    "    opt = ADAM(0.02)\n",
    "    loss(x, y) = crossentropy(logistic_softmax_predict(x, W), onehotbatch(y,0:9))\n",
    "\n",
    "    loss_value = 0.0\n",
    "    epoch_num = 0\n",
    "\n",
    "    #Training loop\n",
    "    for epoch_num in 1:num_epochs\n",
    "        prev_loss_value = loss_value\n",
    "        \n",
    "        #Loop over mini-batches in epoch\n",
    "        start_time = time_ns()\n",
    "        for batch in Iterators.partition(1:n_train, mini_batch_size)\n",
    "            gs = gradient(()->loss(X_total'[:,batch], train_y[batch]), params(W))\n",
    "            for p in (W,)\n",
    "                update!(opt, p, gs[p])\n",
    "            end\n",
    "        end\n",
    "        end_time = time_ns()\n",
    "\n",
    "        #record/display progress\n",
    "        epoch_num += 1\n",
    "        loss_value = loss(X_total', train_y)\n",
    "        println(\"Epoch = $epoch_num ($(round((end_time-start_time)/1e9,digits=2)) sec) Loss = $loss_value\")        \n",
    "    end\n",
    "    return W\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch = 2 (8.7 sec) Loss = 1.7239047373639682\n",
      "Epoch = 3 (5.29 sec) Loss = 1.0180409007626112\n",
      "Epoch = 4 (5.58 sec) Loss = 0.7939230320557332\n",
      "Epoch = 5 (5.6 sec) Loss = 0.6725705514246488\n",
      "Epoch = 6 (5.57 sec) Loss = 0.5947298538518564\n",
      "Epoch = 7 (5.59 sec) Loss = 0.5395838202235905\n",
      "Epoch = 8 (5.7 sec) Loss = 0.49791725265881154\n",
      "Epoch = 9 (6.05 sec) Loss = 0.4653693151230564\n",
      "Epoch = 10 (6.51 sec) Loss = 0.43924988192812864\n",
      "Epoch = 11 (5.84 sec) Loss = 0.41796796020110644\n",
      "Epoch = 12 (5.87 sec) Loss = 0.40029340216611226\n",
      "Epoch = 13 (6.39 sec) Loss = 0.38533789518321715\n",
      "Epoch = 14 (5.7 sec) Loss = 0.372478676681453\n",
      "Epoch = 15 (5.63 sec) Loss = 0.3612562430424544\n",
      "Epoch = 16 (5.27 sec) Loss = 0.35134139457241437\n",
      "Epoch = 17 (5.17 sec) Loss = 0.3425086780990951\n",
      "Epoch = 18 (5.21 sec) Loss = 0.3346090028665252\n",
      "Epoch = 19 (5.27 sec) Loss = 0.3275363932741277\n",
      "Epoch = 20 (5.56 sec) Loss = 0.32119988894055646\n",
      "Epoch = 21 (5.57 sec) Loss = 0.31551023230648045\n",
      "Epoch = 22 (5.56 sec) Loss = 0.3103796746011112\n",
      "Epoch = 23 (5.61 sec) Loss = 0.3057261691139389\n",
      "Epoch = 24 (5.59 sec) Loss = 0.3014758211486228\n",
      "Epoch = 25 (5.64 sec) Loss = 0.2975631973187846\n",
      "Epoch = 26 (5.55 sec) Loss = 0.293931155585798\n",
      "Epoch = 27 (5.54 sec) Loss = 0.2905308350402371\n",
      "Epoch = 28 (5.53 sec) Loss = 0.28732161415245827\n",
      "Epoch = 29 (5.6 sec) Loss = 0.28427084377351164\n",
      "Epoch = 30 (5.53 sec) Loss = 0.28135332716974765\n",
      "Epoch = 31 (5.55 sec) Loss = 0.2785505964044129\n",
      "Epoch = 32 (5.9 sec) Loss = 0.2758500436620855\n",
      "Epoch = 33 (5.61 sec) Loss = 0.273243957632442\n",
      "Epoch = 34 (5.59 sec) Loss = 0.27072851784913066\n",
      "Epoch = 35 (5.31 sec) Loss = 0.2683027614042818\n",
      "Epoch = 36 (5.33 sec) Loss = 0.26596751527837903\n",
      "Epoch = 37 (5.32 sec) Loss = 0.26372440748796266\n",
      "Epoch = 38 (5.31 sec) Loss = 0.26157509888841346\n",
      "Epoch = 39 (5.33 sec) Loss = 0.25952075272804054\n",
      "Epoch = 40 (5.37 sec) Loss = 0.25756170778248233\n",
      "Epoch = 41 (5.31 sec) Loss = 0.2556973290632858\n",
      "Epoch = 42 (5.64 sec) Loss = 0.25392600435603757\n",
      "Epoch = 43 (6.07 sec) Loss = 0.25224524568295054\n",
      "Epoch = 44 (5.36 sec) Loss = 0.2506518534503616\n",
      "Epoch = 45 (5.46 sec) Loss = 0.2491421055828877\n",
      "Epoch = 46 (5.5 sec) Loss = 0.24771194237482494\n",
      "Epoch = 47 (5.54 sec) Loss = 0.24635712842210694\n",
      "Epoch = 48 (5.46 sec) Loss = 0.2450733832865328\n",
      "Epoch = 49 (5.39 sec) Loss = 0.24385648045284744\n",
      "Epoch = 50 (5.28 sec) Loss = 0.24270231895272973\n",
      "Epoch = 51 (5.38 sec) Loss = 0.24160697381205568\n",
      "Epoch = 52 (5.3 sec) Loss = 0.24056673071225673\n",
      "Epoch = 53 (5.31 sec) Loss = 0.23957810806742316\n",
      "Epoch = 54 (5.45 sec) Loss = 0.2386378677419798\n",
      "Epoch = 55 (5.33 sec) Loss = 0.2377430150772292\n",
      "Epoch = 56 (5.34 sec) Loss = 0.2368907896940191\n",
      "Epoch = 57 (5.3 sec) Loss = 0.2360786496510275\n",
      "Epoch = 58 (5.32 sec) Loss = 0.23530425201749106\n",
      "Epoch = 59 (5.3 sec) Loss = 0.23456543256637177\n",
      "Epoch = 60 (5.28 sec) Loss = 0.23386018646975054\n",
      "Epoch = 61 (5.32 sec) Loss = 0.2331866510173527\n",
      "Epoch = 62 (5.33 sec) Loss = 0.23254309072267784\n",
      "Epoch = 63 (5.28 sec) Loss = 0.23192788477554105\n",
      "Epoch = 64 (5.27 sec) Loss = 0.23133951659308785\n",
      "Epoch = 65 (5.33 sec) Loss = 0.23077656514134287\n",
      "Epoch = 66 (5.31 sec) Loss = 0.2302376976884066\n",
      "Epoch = 67 (5.31 sec) Loss = 0.2297216636722226\n",
      "Epoch = 68 (5.3 sec) Loss = 0.22922728940124232\n",
      "Epoch = 69 (5.3 sec) Loss = 0.22875347334701596\n",
      "Epoch = 70 (5.32 sec) Loss = 0.2282991818309286\n",
      "Epoch = 71 (5.3 sec) Loss = 0.22786344495196162\n",
      "Epoch = 72 (5.32 sec) Loss = 0.22744535264714827\n",
      "Epoch = 73 (5.29 sec) Loss = 0.22704405081903511\n",
      "Epoch = 74 (5.4 sec) Loss = 0.2266587375021543\n",
      "Epoch = 75 (5.34 sec) Loss = 0.22628865907076684\n",
      "Epoch = 76 (5.37 sec) Loss = 0.22593310651154858\n",
      "Epoch = 77 (5.3 sec) Loss = 0.2255914117973679\n",
      "Epoch = 78 (5.3 sec) Loss = 0.22526294440292385\n",
      "Epoch = 79 (5.29 sec) Loss = 0.22494710800152604\n",
      "Epoch = 80 (5.75 sec) Loss = 0.22464333737666592\n",
      "Epoch = 81 (5.42 sec) Loss = 0.22435109557404553\n",
      "Epoch = 82 (5.36 sec) Loss = 0.2240698713108985\n",
      "Epoch = 83 (5.3 sec) Loss = 0.223799176650933\n",
      "Epoch = 84 (5.29 sec) Loss = 0.2235385449459495\n",
      "Epoch = 85 (5.4 sec) Loss = 0.22328752903981774\n",
      "Epoch = 86 (5.31 sec) Loss = 0.22304569972744867\n",
      "Epoch = 87 (5.34 sec) Loss = 0.22281264446074292\n",
      "Epoch = 88 (5.33 sec) Loss = 0.22258796629485025\n",
      "Epoch = 89 (5.29 sec) Loss = 0.2223712830704686\n",
      "Epoch = 90 (5.29 sec) Loss = 0.22216222682978837\n",
      "Epoch = 91 (5.33 sec) Loss = 0.22196044346299434\n",
      "Epoch = 92 (5.32 sec) Loss = 0.22176559257679232\n",
      "Epoch = 93 (5.31 sec) Loss = 0.22157734756453246\n",
      "Epoch = 94 (5.32 sec) Loss = 0.22139539583886933\n",
      "Epoch = 95 (5.28 sec) Loss = 0.22121943916454231\n",
      "Epoch = 96 (5.37 sec) Loss = 0.22104919400566317\n",
      "Epoch = 97 (5.3 sec) Loss = 0.22088439178603755\n",
      "Epoch = 98 (5.37 sec) Loss = 0.22072477896014145\n",
      "Epoch = 99 (5.32 sec) Loss = 0.22057011681140415\n",
      "Epoch = 100 (5.42 sec) Loss = 0.22042018093287757\n",
      "Epoch = 101 (5.33 sec) Loss = 0.2202747603959889\n",
      "Epoch = 102 (5.31 sec) Loss = 0.22013365666379253\n",
      "Epoch = 103 (5.3 sec) Loss = 0.2199966823432611\n",
      "Epoch = 104 (5.32 sec) Loss = 0.21986365988809467\n",
      "Epoch = 105 (5.31 sec) Loss = 0.21973442035764315\n",
      "Epoch = 106 (5.28 sec) Loss = 0.21960880231384192\n",
      "Epoch = 107 (5.37 sec) Loss = 0.2194866509053386\n",
      "Epoch = 108 (5.29 sec) Loss = 0.2193678171552072\n",
      "Epoch = 109 (5.33 sec) Loss = 0.21925215744221027\n",
      "Epoch = 110 (5.3 sec) Loss = 0.2191395331483837\n",
      "Epoch = 111 (5.35 sec) Loss = 0.2190298104376532\n",
      "Epoch = 112 (5.3 sec) Loss = 0.21892286012928908\n",
      "Epoch = 113 (5.34 sec) Loss = 0.21881855763370106\n",
      "Epoch = 114 (5.29 sec) Loss = 0.21871678292404043\n",
      "Epoch = 115 (5.3 sec) Loss = 0.2186174205235739\n",
      "Epoch = 116 (5.53 sec) Loss = 0.21852035949479706\n",
      "Epoch = 117 (5.36 sec) Loss = 0.21842549342123904\n",
      "Epoch = 118 (5.33 sec) Loss = 0.21833272037675952\n",
      "Epoch = 119 (5.49 sec) Loss = 0.21824194287992876\n",
      "Epoch = 120 (5.45 sec) Loss = 0.21815306783298782\n",
      "Epoch = 121 (5.37 sec) Loss = 0.2180660064461097\n",
      "Epoch = 122 (5.33 sec) Loss = 0.21798067414840408\n",
      "Epoch = 123 (5.33 sec) Loss = 0.21789699048748012\n",
      "Epoch = 124 (5.27 sec) Loss = 0.2178148790195232\n",
      "Epoch = 125 (5.29 sec) Loss = 0.21773426719182692\n",
      "Epoch = 126 (5.31 sec) Loss = 0.21765508621962823\n",
      "Epoch = 127 (5.31 sec) Loss = 0.2175772709589381\n",
      "Epoch = 128 (5.3 sec) Loss = 0.21750075977688835\n",
      "Epoch = 129 (5.31 sec) Loss = 0.21742549442092277\n",
      "Epoch = 130 (5.38 sec) Loss = 0.2173514198879707\n",
      "Epoch = 131 (5.33 sec) Loss = 0.2172784842945398\n",
      "Epoch = 132 (5.34 sec) Loss = 0.2172066387484659\n",
      "Epoch = 133 (5.34 sec) Loss = 0.21713583722285076\n",
      "Epoch = 134 (5.36 sec) Loss = 0.21706603643251046\n",
      "Epoch = 135 (5.35 sec) Loss = 0.21699719571304704\n",
      "Epoch = 136 (5.28 sec) Loss = 0.2169292769024591\n",
      "Epoch = 137 (5.46 sec) Loss = 0.21686224422503286\n",
      "Epoch = 138 (6.59 sec) Loss = 0.21679606417712485\n",
      "Epoch = 139 (6.26 sec) Loss = 0.21673070541439332\n",
      "Epoch = 140 (5.9 sec) Loss = 0.21666613864006226\n",
      "Epoch = 141 (5.93 sec) Loss = 0.21660233649394753\n",
      "Epoch = 142 (6.31 sec) Loss = 0.2165392734422138\n",
      "Epoch = 143 (6.25 sec) Loss = 0.21647692566815538\n",
      "Epoch = 144 (6.28 sec) Loss = 0.2164152709646455\n",
      "Epoch = 145 (6.48 sec) Loss = 0.2163542886292091\n",
      "Epoch = 146 (6.4 sec) Loss = 0.21629395936287796\n",
      "Epoch = 147 (6.34 sec) Loss = 0.21623426517401906\n",
      "Epoch = 148 (6.42 sec) Loss = 0.21617518928817905\n",
      "Epoch = 149 (6.5 sec) Loss = 0.2161167160646715\n",
      "Epoch = 150 (6.36 sec) Loss = 0.21605883092022168\n",
      "Epoch = 151 (6.58 sec) Loss = 0.21600152025954883\n",
      "Epoch = 152 (7.16 sec) Loss = 0.21594477141239204\n",
      "Epoch = 153 (7.99 sec) Loss = 0.2158885725762238\n",
      "Epoch = 154 (7.86 sec) Loss = 0.21583291276376804\n",
      "Epoch = 155 (9.0 sec) Loss = 0.2157777817544329\n",
      "Epoch = 156 (6.32 sec) Loss = 0.21572317004885874\n",
      "Epoch = 157 (6.3 sec) Loss = 0.21566906882592485\n",
      "Epoch = 158 (6.28 sec) Loss = 0.21561546990171956\n",
      "Epoch = 159 (8.0 sec) Loss = 0.21556236569013412\n",
      "Epoch = 160 (7.96 sec) Loss = 0.21550974916487098\n",
      "Epoch = 161 (6.76 sec) Loss = 0.21545761382275552\n",
      "Epoch = 162 (7.5 sec) Loss = 0.21540595364830978\n",
      "Epoch = 163 (6.67 sec) Loss = 0.2153547630795865\n",
      "Epoch = 164 (6.9 sec) Loss = 0.2153040369752846\n",
      "Epoch = 165 (6.69 sec) Loss = 0.2152537705831681\n",
      "Epoch = 166 (6.54 sec) Loss = 0.21520395950981075\n",
      "Epoch = 167 (6.83 sec) Loss = 0.21515459969167575\n",
      "Epoch = 168 (9.05 sec) Loss = 0.21510568736752905\n",
      "Epoch = 169 (6.9 sec) Loss = 0.21505721905217434\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch = 170 (6.81 sec) Loss = 0.21500919151148568\n",
      "Epoch = 171 (6.83 sec) Loss = 0.21496160173870624\n",
      "Epoch = 172 (6.73 sec) Loss = 0.21491444693197687\n",
      "Epoch = 173 (6.72 sec) Loss = 0.2148677244730512\n",
      "Epoch = 174 (8.18 sec) Loss = 0.2148214319071547\n",
      "Epoch = 175 (7.51 sec) Loss = 0.21477556692394303\n",
      "Epoch = 176 (7.18 sec) Loss = 0.21473012733951616\n",
      "Epoch = 177 (7.31 sec) Loss = 0.21468511107944505\n",
      "Epoch = 178 (7.13 sec) Loss = 0.21464051616277194\n",
      "Epoch = 179 (6.63 sec) Loss = 0.2145963406869443\n",
      "Epoch = 180 (7.46 sec) Loss = 0.214552582813649\n",
      "Epoch = 181 (6.68 sec) Loss = 0.2145092407555109\n",
      "Epoch = 182 (7.87 sec) Loss = 0.21446631276362743\n",
      "Epoch = 183 (7.31 sec) Loss = 0.21442379711590798\n",
      "Epoch = 184 (7.93 sec) Loss = 0.21438169210619284\n",
      "Epoch = 185 (7.57 sec) Loss = 0.21433999603412526\n",
      "Epoch = 186 (7.33 sec) Loss = 0.21429870719575267\n",
      "Epoch = 187 (7.36 sec) Loss = 0.21425782387483464\n",
      "Epoch = 188 (6.99 sec) Loss = 0.21421734433483514\n",
      "Epoch = 189 (6.65 sec) Loss = 0.21417726681157873\n",
      "Epoch = 190 (7.34 sec) Loss = 0.21413758950654907\n",
      "Epoch = 191 (7.57 sec) Loss = 0.21409831058081144\n",
      "Epoch = 192 (8.06 sec) Loss = 0.214059428149539\n",
      "Epoch = 193 (7.09 sec) Loss = 0.21402094027712462\n",
      "Epoch = 194 (7.1 sec) Loss = 0.2139828449728593\n",
      "Epoch = 195 (7.41 sec) Loss = 0.21394514018716065\n",
      "Epoch = 196 (7.6 sec) Loss = 0.213907823808332\n",
      "Epoch = 197 (8.01 sec) Loss = 0.21387089365983764\n",
      "Epoch = 198 (6.78 sec) Loss = 0.21383434749807487\n",
      "Epoch = 199 (7.56 sec) Loss = 0.21379818301062847\n",
      "Epoch = 200 (8.38 sec) Loss = 0.21376239781499012\n",
      "Epoch = 201 (7.62 sec) Loss = 0.21372698945772695\n"
     ]
    }
   ],
   "source": [
    "W = train_softmax_logistic();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "accuracy_softmax = 0.9221\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "10×10 Matrix{Int64}:\n",
       " 944     0   18    4    0   23   18    5   14   15\n",
       "   0  1107   54   17   22   18   10   40   46   11\n",
       "   1     2  813   23    6    3    9   16   11    2\n",
       "   2     2   26  880    1   72    0    6   30   17\n",
       "   2     3   15    5  881   24   22   26   27   80\n",
       "   7     1    0   17    5  659   17    0   40    1\n",
       "  14     5   42    9   10   23  875    1   15    1\n",
       "   2     1   22   21    2   14    0  884   12   77\n",
       "   7    14   37   22   11   39    7    0  759    4\n",
       "   1     0    5   12   44   17    0   50   20  801"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predict_softmax(img) = argmax(logistic_softmax_predict(vcat(1,vec(img)),W)) - 1\n",
    "accuracy_softmax = mean([predict_softmax(test_x[:,:,i]) == test_y[i] for i in 1:10000])\n",
    "\n",
    "@show accuracy_softmax\n",
    "conf_matrix(predict_softmax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.10.0",
   "language": "julia",
   "name": "julia-1.10"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
